#include "png/png.h"
#include "png/pngstruct.h"
#include "png/pnginfo.h"
#include "png/pngpriv.h"

#include "Image.h"
#include "Resource.h"

namespace Resource
{
	static Image* loadImagePNG(const char *name);
	
	Image* loadImage(const char *name)
	{
		ASSERT(name);

		const char *ext = strrchr(name, '.');
		if (!ext)
		{
			LOG_ERROR("Invalid image '%s' extension\n", name);
			return NULL;
		}

		if (strcmp(ext, ".png") == 0)
			return loadImagePNG(name);

		LOG_ERROR("Not supported image '%s' extension '%s'\n", name, ext);
		return NULL;
	}
	
	static void PNGAPI pngErrorCallback(png_structp png_ptr, png_const_charp msg)
	{
		LOG_ERROR("Failed to read PNG: %s\n", msg);
		longjmp(png_ptr->longjmp_buffer, 1);
	}
	
	static void PNGAPI pngRead(png_structp png_ptr, png_bytep data, png_size_t length)
	{
		ASSERT(png_ptr->io_ptr);
		memcpy(data, png_ptr->io_ptr, length);
		png_ptr->io_ptr = (uint8_t *)png_ptr->io_ptr + length;
	}
	
	static Image* loadImagePNG(const char *name)
	{
		ASSERT(name);
		
		File *file = loadFile(name);
		if (file == NULL)
			return NULL;
		
		uint32_t     bpp;
		png_structp  png_ptr;
		png_infop    info_ptr;
		png_uint_32  width, height;
		int32_t      bit_depth, color_type, stride;
		png_bytep   *row_pointers;
		
		if (png_sig_cmp((png_byte *)file->data, 0, 8))
		{
			LOG_ERROR("Wrong PNG format '%s'\n", name);
			freeFile(file);
			return NULL;
		}
		
		if ((png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING,
			NULL, pngErrorCallback, NULL)) == NULL)
		{
			LOG_ERROR("Failed to create PNG read struct\n");
			freeFile(file);
			return NULL;
		}
		
		if ((info_ptr = png_create_info_struct(png_ptr)) == NULL)
		{
			LOG_ERROR("Failed to create PNG info struct\n");
			png_destroy_read_struct(&png_ptr, NULL, NULL);
			freeFile(file);
			return NULL;
		}
		
		if (setjmp(png_jmpbuf(png_ptr)))
		{
			png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
			freeFile(file);
			return NULL;
		}
		
		png_set_read_fn(png_ptr, (void *)(file->data + 8), pngRead);
		png_set_sig_bytes(png_ptr, 8);
		png_read_info(png_ptr, info_ptr);
		
		png_get_IHDR(png_ptr, info_ptr, &width, &height, &bit_depth, &color_type, NULL, NULL, NULL);
		
		png_set_strip_16(png_ptr);
		png_set_packing(png_ptr);
		
		if (color_type == PNG_COLOR_TYPE_PALETTE)
			png_set_palette_to_rgb(png_ptr);
		
		if (color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8)
			png_set_expand_gray_1_2_4_to_8(png_ptr);
		
		if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
			png_set_tRNS_to_alpha(png_ptr);
		
		//if (color_type == PNG_COLOR_TYPE_GRAY || color_type == PNG_COLOR_TYPE_GRAY_ALPHA)
		//	png_set_gray_to_rgb(png_ptr);
		
		png_read_update_info(png_ptr, info_ptr);
		
		png_get_IHDR(png_ptr, info_ptr, &width, &height, &bit_depth, &color_type, NULL, NULL, NULL);
		
		switch (color_type)
		{
			case PNG_COLOR_TYPE_GRAY:
			case PNG_COLOR_TYPE_GRAY_ALPHA:
				bpp = 8;
			break;
			case PNG_COLOR_TYPE_RGB:
				bpp = 24;
			break;
			case PNG_COLOR_TYPE_RGBA:
				bpp = 32;
			break;
			default:
				LOG_ERROR("Unknown PNG color type %d in '%s'\n", color_type, name);
				png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
				freeFile(file);
				return NULL;
		}
		
		stride = png_get_rowbytes(png_ptr, info_ptr);
		
		row_pointers = new png_bytep[height];
		ASSERT(row_pointers);
		
		Image *image = new Image(width, height, bpp);
		
		// mirror image vertically
		for (uint32_t i = 0; i < height; ++i)
			row_pointers[i] = (png_bytep)image->data + (height - i - 1) * stride;
		
		png_read_image(png_ptr, row_pointers);
		png_read_end(png_ptr, info_ptr);
		png_destroy_read_struct(&png_ptr, &info_ptr, NULL);
		
		delete[] row_pointers;
		freeFile(file);
		
		return image;
	}
}
